package com.sz.admin.ai.factory.chatFactory.handler.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.GenerationUsage;
import com.alibaba.dashscope.common.Message;
import com.mybatisflex.core.query.QueryWrapper;
import com.sz.admin.ai.factory.chatFactory.entity.ChatPlatformEnum;
import com.sz.admin.ai.factory.chatFactory.entity.MessageTypoeEnum;
import com.sz.admin.ai.factory.chatFactory.entity.dto.ChatDTO;
import com.sz.admin.ai.factory.chatFactory.entity.dto.ChatEmbeddingModelDTO;
import com.sz.admin.ai.factory.chatFactory.entity.dto.ChatMessageDTO;
import com.sz.admin.ai.factory.chatFactory.entity.dto.ChatModelDTO;
import com.sz.admin.ai.factory.chatFactory.handler.ChatHandler;
import com.sz.admin.ai.factory.embeddingFactory.EmbeddingHandlerFactory;
import com.sz.admin.ai.factory.embeddingFactory.Handler.EmbeddingHandler;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingHandlerRequestDTO;
import com.sz.admin.aichatmessage.pojo.dto.AiChatMessageUpdateDTO;
import com.sz.admin.aichatmessage.pojo.po.AiChatMessage;
import com.sz.admin.aichatmessage.service.AiChatMessageService;
import com.sz.core.util.JsonUtils;
import com.sz.core.util.StringUtils;
import io.reactivex.Flowable;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.log4j.Log4j2;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @描述:
 * @Author: TangYuan
 * @Date: 2025/3/7 16:26
 */
@Component
@AllArgsConstructor
@Log4j2
public class AliyunThinkChatHandlerImpl implements ChatHandler {
    
    private final AiChatMessageService aiChatMessageService;
    
    /**
     * 是否使用文件问答
     * @param promptUserSpec promptUserSpec
     * @param chatMessageDTO chatMessageDTO
     * @param contextNumber
     */
    private List<Message> useFileChat(ChatMessageDTO chatMessageDTO, String contextNumber) {
        List<Message> messageList = new ArrayList<>();
        
        List<AiChatMessage> list = aiChatMessageService.list(QueryWrapper.create()
                .eq("session_id", chatMessageDTO.getSessionId())
                .eq("type", MessageTypoeEnum.USER.getCapitalString())
                .ne("total_token", 0)
                .orderBy("id", false)
                .limit(Integer.valueOf(contextNumber)));
        Collections.reverse(list);
        list.forEach(aiChatMessage -> {
            Message userMsg = Message.builder()
                    .role(MessageTypoeEnum.getEnum(aiChatMessage.getType()).getLowercaseString())
                    .content(aiChatMessage.getContent())
                    .build();
            messageList.add(userMsg);
        });
        
        Message userMsg = Message.builder()
                .role(MessageTypoeEnum.getEnum(chatMessageDTO.getRole()).getLowercaseString())
                .content(chatMessageDTO.getContent())
                .build();
        messageList.add(userMsg);
        
        messageList.forEach(message -> {
            System.out.println(message.getRole() + ":" + message.getContent());
        });
        return messageList;
    }
    
    /**
     * 是否使用向量数据库
     * @param chatMessageDTO          chatMessageDTO
     * @param embeddingMode           embeddingMode
     * @param advisorSpec             advisorSpec
     * @param embeddingHandlerFactory embeddingHandlerFactory
     */
    private static void useEmbedding(ChatMessageDTO chatMessageDTO, ChatEmbeddingModelDTO embeddingMode, ChatClient.AdvisorSpec advisorSpec, EmbeddingHandlerFactory embeddingHandlerFactory) {
        if (embeddingMode != null) {
            EmbeddingHandler ollamaEmbedding = embeddingHandlerFactory.getHandler(embeddingMode.getModelPlatformName());
            // 是否使用向量数据库
            if (ollamaEmbedding != null && embeddingMode.getKnowledgeBaseId() != null) {
                EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO = EmbeddingHandlerRequestDTO.builder()
                        // 设置产品平台
                        .thirdPartyName(embeddingMode.getModelPlatformName())
                        .modelName(embeddingMode.getModelName())
                        .modelKey("")
                        .userId(embeddingMode.getUserId())
                        .knowledgeBaseId(embeddingMode.getKnowledgeBaseId())
                        .build();
                ollamaEmbedding.useVectorStore(embeddingHandlerRequestDTO, 5, advisorSpec, chatMessageDTO.getContent());
            }
        }
    }
    
    /**
     * 开始时，将消息存进数据库
     * @param chat chat
     * @return {@link AiChatMessage}
     */
    private AiChatMessage startMessage(ChatDTO chat) {
        // 保存用户发送的消息
        AiChatMessage aiChatMessage = new AiChatMessage();
        aiChatMessage.setThirdPartyName(chat.getUser().getThirdPartyName());
        aiChatMessage.setSessionId(chat.getMessage().getSessionId());
        aiChatMessage.setType(MessageTypoeEnum.getEnum(chat.getMessage().getRole()).getCapitalString());
        aiChatMessage.setContent(chat.getMessage().getContent());
        aiChatMessage.setPlatformName(chat.getChatModel().getModelPlatformName());
        aiChatMessage.setUseModel(chat.getChatModel().getModelName());
        aiChatMessage.setTotalToken(0);
        aiChatMessageService.save(aiChatMessage);
        return aiChatMessage;
    }
    
    /**
     * 结束时，将消息存进数据库，并修改请求时的token
     */
    private void endMessage(ChatDTO chat, AiChatMessage aiChatMessage, GenerationUsage usage, String chatMessage, String chatThinkMessage, String capitalString) {
        // 更新用户发送的消息使用的token
        AiChatMessageUpdateDTO aiChatMessageUpdateDTO = new AiChatMessageUpdateDTO();
        aiChatMessageUpdateDTO.setId(aiChatMessage.getId());
        aiChatMessageUpdateDTO.setTotalToken(usage.getInputTokens());
        aiChatMessageService.update(aiChatMessageUpdateDTO);
        
        // 保存助手返回的消息记录
        AiChatMessage aiChatMessage1 = new AiChatMessage();
        aiChatMessage1.setThirdPartyName(chat.getUser().getThirdPartyName());
        aiChatMessage1.setSessionId(chat.getMessage().getSessionId());
        aiChatMessage1.setType(MessageTypoeEnum.getEnum(capitalString).getCapitalString());
        aiChatMessage1.setContent(chatMessage);
        aiChatMessage1.setReasoningContent(chatThinkMessage);
        aiChatMessage1.setPlatformName(chat.getChatModel().getModelPlatformName());
        aiChatMessage1.setUseModel(chat.getChatModel().getModelName());
        aiChatMessage1.setTotalToken(usage.getOutputTokens());
        aiChatMessageService.save(aiChatMessage1);
    }
    
    /**
     * 获取ChatPlatformEnum
     * @return {@link ChatPlatformEnum}
     */
    @Override
    public ChatPlatformEnum getChatPlatformEnum() {
        return ChatPlatformEnum.ALIYUNTHINK;
    }
    
    @SneakyThrows
    @Override
    public Flux<ServerSentEvent<String>> chat(ChatDTO chat) {
        
        ChatMessageDTO chatMessageDTO = chat.getMessage();
        ChatModelDTO chatModelDTO = chat.getChatModel();
        
        // 开始保存发送信息
        AiChatMessage aiChatMessage = startMessage(chat);
        
        Generation gen = new Generation();
        
        GenerationParam param = GenerationParam.builder()
                .apiKey("sk-1ee02e3b0e564d18b289e818f485f763")
                .model(chatModelDTO.getModelName())
                .prompt(chat.getMessage().getPromptText())
                .messages(useFileChat(chatMessageDTO, chatModelDTO.getContextNumber()))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .temperature(StringUtils.isEmpty(chatModelDTO.getTemperature()) ? 0.7F : Float.parseFloat(chatModelDTO.getTemperature()))
                .build();
        
        Flowable<GenerationResult> generationResultFlowable = gen.streamCall(param);
        
        StringBuilder chatMessage = new StringBuilder();
        StringBuilder thinkMessage = new StringBuilder();
        Flowable<ServerSentEvent<String>> message = generationResultFlowable.map(chatResponse -> {
            // 拼接流式输出
            chatMessage.append(chatResponse.getOutput().getChoices().getFirst().getMessage().getContent());
            thinkMessage.append(chatResponse.getOutput().getChoices().getFirst().getMessage().getReasoningContent());
            
            // 结束保存发送信息
            if (chatResponse.getOutput().getChoices().getFirst().getFinishReason().equals("STOP") || chatResponse.getOutput().getChoices().getFirst().getFinishReason().equals("stop")) {
                endMessage(chat, aiChatMessage, chatResponse.getUsage(), String.valueOf(chatMessage), String.valueOf(thinkMessage), chatResponse.getOutput().getChoices().getFirst().getMessage().getRole());
            }
            // 前端返回
            return ServerSentEvent
                    .builder(JsonUtils.toJsonString(chatResponse))
                    .event("message")
                    .build();
        });
        return Flux.from(message);
    }
}